package cn.doitedu.datayi.utils

import cn.doitedu.datayi.utils.BitmapUDAF.{desBitmap, serBitmap}
import org.apache.spark.sql.{Encoder, Encoders}
import org.apache.spark.sql.expressions.Aggregator
import org.roaringbitmap.RoaringBitmap

object BitMapAggr extends Aggregator[Array[Byte],Array[Byte],Array[Byte]]{
  override def zero: Array[Byte] = {
    val buffBitmap = RoaringBitmap.bitmapOf()
    // 序列化
    serBitmap(buffBitmap)
  }

  override def reduce(b: Array[Byte], a: Array[Byte]): Array[Byte] = {
    val bm_b = desBitmap(b)
    val bm_a = desBitmap(a)
    bm_b.or(bm_a)

    serBitmap(bm_b)
  }

  override def merge(b1: Array[Byte], b2: Array[Byte]): Array[Byte] = reduce(b1,b2)

  override def finish(reduction: Array[Byte]): Array[Byte] = reduction

  override def bufferEncoder: Encoder[Array[Byte]] = Encoders.BINARY

  override def outputEncoder: Encoder[Array[Byte]] = Encoders.BINARY
}
